modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py (1)

137-147: Clarify engine_path semantics in the docstring

Doc says “Path to save the TensorRT engine,” but the code (and suggested change) preserves that path and also writes a hashed copy into output_dir. Make this explicit.

-        engine_path: Path to save the TensorRT engine.
+        engine_path: Optional path where trtexec will initially write the engine (via --saveEngine).
+            The final, hashed artifact is stored under `output_dir` (or the default artifact dir)
+            as `<hash>-{model_name}.engine`. When a custom `engine_path` is provided, that file
+            is preserved and a hashed copy is saved in the artifact dir.

modelopt/torch/_deploy/_runtime/trt_client.py (1)

77-79: Doc: reference builder’s engine_path behavior

Make it clear this is forwarded to the builder and see builder docs for final file placement.

examples/onnx_ptq/evaluate.py (1)

38-43: Clarify that --engine_path is an output path used during build

Avoid confusion with a “load existing engine” path.

-    parser.add_argument(
+    parser.add_argument(
         "--engine_path",
         type=str,
         required=True,
-        help="Path to the TensorRT engine",
+        help="Output path where trtexec should write the engine (will also be copied to the artifact dir)",
     )

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

tests/examples/test_onnx_ptq.sh (1)

203-214: Fix array membership test (ShellCheck SC2199/SC2076).

Use glob match on the space-joined array; current regex-with-quotes is brittle and flagged by ShellCheck.

Apply this diff:

-        if [[ " ${latency_models[@]} " =~ " $model_name " ]]; then
+        # Use glob match on joined array to test membership (SC2199/SC2076)
+        if [[ " ${latency_models[*]} " == *" $model_name "* ]]; then
             CUDA_VISIBLE_DEVICES=$gpu_id python evaluate.py \
                 --onnx_path=$eval_model_path \
-                --engine_path=$engine_path \
+                --engine_path=$engine_path \
                 --model_name="${timm_model_name[$model_name]}" \
                 --engine_precision=$precision \
                 --results_path=$model_dir/$quant_mode/${model_name}_${quant_mode}.csv &

tests/examples/test_onnx_ptq.sh (7)

38-58: Harden CLI argument parsing (shift inside for-loop is ineffective).

Switch to a while/case parser so shifts take effect and unknown flags are surfaced.

Apply this diff:

-# Parse arguments
-clean_mode=true
-imagenet_path=""
-models_folder=""
-
-for arg in "$@"; do
-    case $arg in
-        --no-clean)
-            clean_mode=false
-            shift
-            ;;
-        *)
-            if [ -z "$imagenet_path" ]; then
-                imagenet_path="$arg"
-            elif [ -z "$models_folder" ]; then
-                models_folder="$arg"
-            fi
-            shift
-            ;;
-    esac
-done
+# Parse arguments
+clean_mode=true
+imagenet_path=""
+models_folder=""
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --no-clean)
+            clean_mode=false
+            shift
+            ;;
+        -*)
+            echo "Unknown option: $1" >&2
+            exit 2
+            ;;
+        *)
+            if [[ -z "$imagenet_path" ]]; then
+                imagenet_path="$1"
+            elif [[ -z "$models_folder" ]]; then
+                models_folder="$1"
+            else
+                echo "Unexpected positional arg: $1" >&2
+                exit 2
+            fi
+            shift
+            ;;
+    esac
+done

---

81-86: Make model discovery robust when no .onnx files exist.

Enable nullglob to avoid iterating a literal pattern and fail fast with a clear message.

Apply this diff:

-declare -a model_paths
-for model in "$models_folder"/*.onnx; do
-    model_paths+=("$model")
-done
+shopt -s nullglob
+declare -a model_paths=("$models_folder"/*.onnx)
+shopt -u nullglob
+if [[ ${#model_paths[@]} -eq 0 ]]; then
+    echo "ERROR: No .onnx models found in: $models_folder" >&2
+    exit 1
+fi

---

64-66: Validate input directories early with helpful errors.

Catch misconfigured imagenet/models paths before doing work.

Apply this diff:

 imagenet_path=${imagenet_path:-/data/imagenet/}
 models_folder=${models_folder:-/models/onnx}
+if [[ ! -d "$imagenet_path" ]]; then
+  echo "ERROR: imagenet_path not found: $imagenet_path" >&2
+  exit 1
+fi
+if [[ ! -d "$models_folder" ]]; then
+  echo "ERROR: models_folder not found: $models_folder" >&2
+  exit 1
+fi

---

36-36: Quote paths and add a trap to always unwind directory stack.

Avoid word-splitting and ensure popd even on early exits.

Apply this diff:

-pushd $public_example_dir
+pushd "$public_example_dir"
+# Always return to original dir, even on errors
+trap 'popd >/dev/null' EXIT

Optionally remove the explicit popd calls at Lines 246 and 250 since the trap handles it.

---

117-127: Quote path expansions to handle spaces safely.

Low-risk now, but quoting prevents subtle bugs.

Apply this diff:

-    model_dir=build/$model_name
+    model_dir="build/$model_name"
@@
-        mkdir -p $model_dir/$mode
+        mkdir -p "$model_dir/$mode"
@@
-    cp $model_path $model_dir/fp16/model.onnx
+    cp "$model_path" "$model_dir/fp16/model.onnx"

---

151-157: Quote CLI args and ensure consistent style.

Keep style consistent (either use --arg=value everywhere or split args uniformly) and quote paths.

Apply this diff:

-        CUDA_VISIBLE_DEVICES=$gpu_id python -m modelopt.onnx.quantization \
-            --onnx_path=$model_dir/fp16/model.onnx \
-            --quantize_mode=$quant_mode \
-            --calibration_data=$calib_data_path \
-            --output_path=$model_dir/$quant_mode/model.quant.onnx \
-            --calibration_eps=cuda:0 &
+        CUDA_VISIBLE_DEVICES=$gpu_id python -m modelopt.onnx.quantization \
+            --onnx_path="$model_dir/fp16/model.onnx" \
+            --quantize_mode="$quant_mode" \
+            --calibration_data="$calib_data_path" \
+            --output_path="$model_dir/$quant_mode/model.quant.onnx" \
+            --calibration_eps="cuda:0" &

Apply similar quoting for evaluate.py invocations.

---

240-248: Optional: gate cleanup symmetry via flag only (keep initial artifacts when troubleshooting).

Current flow always rm -rf build at start; --no-clean only preserves end artifacts. If you'd like to keep prior runs for diffing, consider skipping the initial rm when --no-clean is set.

Example tweak:

-rm -rf build
-mkdir -p build
+if [[ "$clean_mode" == true ]]; then
+  rm -rf build
+fi
+mkdir -p build

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro